import networkx as nx
import matplotlib.pyplot as plt
from collections import deque
from pulp import *
from tools import generate_gas_network, visualize
from matplotlib.font_manager import FontProperties
import pulp as pulp
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径

def calculate_users(G):
    valves = [n for n, attr in G.nodes(data=True) if attr['type'] == 'valve']
    users = 0
    queue = deque(valves)
    visited = set(valves)
    while queue:
        node = queue.popleft()
        if G.nodes[node].get("user_weight"):
            users += G.nodes[node]["user_weight"]
        for neighbor in G.successors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return users

def calculate_leaks(G):
    valves = [n for n, attr in G.nodes(data=True) if attr['type'] == 'valve']
    leaks = 0
    queue = deque(valves)
    visited = set(valves)
    while queue:
        node = queue.popleft()
        if G.nodes[node].get("leak_weight"):
            leaks += G.nodes[node]["leak_weight"]
        for neighbor in G.successors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    return leaks

if __name__ == "__main__":
    show = False
    #生成图
    G0, pos = generate_gas_network()

    if show:
        nx.draw_networkx(G0, pos)
        plt.show()

    #predecessors
    #successors
    #generate_cost_flow_network(G)

    #节点
    leaks = [n for n, attr in G0.nodes(data=True) if attr['type'] == 'leak']
    #sources = [n for n, attr in G0.nodes(data=True) if attr['type'] == 'source']
    valves = [n for n, attr in G0.nodes(data=True) if attr['type'] == 'valve']
    users = [n for n, attr in G0.nodes(data=True) if attr['type'] == 'user']
    M = len(G0.nodes())*10+len(G0.edges())*10
    epsilon = 0.000001

    #加入超级源和超级汇
    G2 = G0.copy()
    for u, v in G2.edges():
        G2[u][v]['capacity'] = M
        G2[u][v]['weight'] = 0
    super_source = "super_source"
    super_sink = "super_sink"
    pos[super_source] = [0,0.3]
    pos[super_sink] = [0,-0.3]
    nodes = []
    nodes.append((super_source, {'type': super_source}))
    nodes.append((super_sink, {'type': super_sink}))
    G2.add_nodes_from(nodes)
    for u in valves:
        G2.add_edge(super_source, u, capacity = M, weight = 0)
    for u in leaks:
        G2.add_edge(u, super_sink, capacity = 1, weight = 0)
    if show:
        nx.draw_networkx(G2, pos)

    #拆点（user节点）
    G3 = G2.copy()
    for u in G2.nodes():
        if G2.nodes[u]['type'] == 'user':
            u_out = f"{u}_out"
            G3.add_nodes_from([u_out])
            for v in G2.successors(u):
                G3.remove_edge(u, v)
                G3.add_edge(u_out, v, capacity = M, weight = 0)
            pos[u_out] = pos[u] + [0.02, 0.02]
            G3.add_edge(u, u_out, capacity = M, weight = 1)
    if show:
        nx.draw_networkx(G3, pos)
        plt.show()

    #给边编号
    idx = 0
    capacity = []
    weight = []
    node_to_edge = {}
    edge_to_node = {}
    equal_edge = []
    leaks_edges = []
    all_edges = []
    for u in G3.nodes:
        for v in G3.successors(u):
            node_to_edge[(u,v)] = idx
            edge_to_node[idx] = (u, v)
            capacity.append(G3[u][v]['capacity'])
            weight.append(G3[u][v]['weight'])
            if u in leaks and v == super_sink:
                leaks_edges.append(idx)
            all_edges.append(idx)
            idx += 1
    for u in G3.nodes:
        out_edge = []
        for v in G3.successors(u):
            out_edge.append(node_to_edge[(u,v)])
        in_edge = []
        for v in G3.predecessors(u):
            in_edge.append(node_to_edge[(v,u)])
        if len(out_edge) > 0 and len(in_edge) > 0:
            equal_edge.append((in_edge, out_edge))

    #线性规划
    prob = pulp.LpProblem("MultiObjective_ILP", pulp.LpMaximize)
    x = pulp.LpVariable.dicts("x", range(idx), lowBound=0, upBound=M, cat=pulp.LpInteger)
    y = pulp.LpVariable.dicts("y", range(idx), cat=pulp.LpBinary)
    big_weight = M * 10  
    small_weight = 1     
    prob += big_weight * pulp.lpSum([x[i] for i in leaks_edges]) \
        - small_weight * pulp.lpSum([weight[i] * y[i] for i in all_edges])
    for in_edge, out_edge in equal_edge:
        prob += pulp.lpSum([x[i] for i in in_edge]) == pulp.lpSum([x[i] for i in out_edge])
    for i in all_edges:
        prob += x[i] <= capacity[i]
    for i in all_edges:
        prob += x[i] <= M * y[i]
        prob += x[i] >= epsilon * y[i]
    prob.solve(pulp.PULP_CBC_CMD())
    print("Status:", pulp.LpStatus[prob.status])
    print("Objective value =", pulp.value(prob.objective))
    x_result = [pulp.value(x[i]) for i in all_edges]
    y_result = [pulp.value(y[i]) for i in all_edges]
    for i in all_edges:
        print (i, edge_to_node[i], y_result[i], x_result[i], weight[i], G3[edge_to_node[i][0]][edge_to_node[i][1]])
    print ("阻止泄漏点：", sum(x_result[i] for i in leaks_edges))
    print ("影响用户数：", sum(weight[i] * y_result[i] for i in all_edges))
    #exit()
    
    #用线性规划求得的流量赋值网络容量，求得最小割
    G4 = G0.copy()
    super_source = "super_source"
    super_sink = "super_sink"
    nodes = []
    nodes.append((super_source, {'type': super_source}))
    nodes.append((super_sink, {'type': super_sink}))
    G4.add_nodes_from(nodes)
    for u in valves:
        G4.add_edge(super_source, u)
    for  u in leaks:
        G4.add_edge(u, super_sink)
    for u, v in G4.edges():
        G4[u][v]['capacity'] = 1e9        
    for i in all_edges:
        u, v = edge_to_node[i]
        if u in valves:
            G4[u][v]['capacity'] = x_result[i]
    maxflow, (reach_set, non_reach_set) = nx.minimum_cut(G4, super_source, super_sink)
    print ("maxflow", maxflow)
    valve_cut = []
    for u in reach_set:
        for v in non_reach_set:
            if (u, v) in G4.edges() and u in valves:
                valve_cut.append(u)
    print ("阀门关闭", valve_cut)
    print ("阀门关闭数", len(valve_cut))
    G5 = G0.copy()
    for u in G0.nodes():
        if u in valve_cut:
            for v in G0.successors(u):
                G5.remove_edge(u, v)
            G5.remove_node(u)
    leaks_ori = calculate_leaks(G0)
    leaks_cut = calculate_leaks(G5)
    users_ori = calculate_users(G0)
    users_cut = calculate_users(G5)
    print ("阻止泄漏点：", leaks_ori, "-", leaks_cut, "=", leaks_ori - leaks_cut)
    print ("影响用户数：", users_ori, "-", users_cut, "=", users_ori - users_cut)
    visualize(G0, pos, valve_cut)
    